#ifndef PCPS_ACCUMULATOR_GRADIENT_HEADER
#define PCPS_ACCUMULATOR_GRADIENT_HEADER

#include <src/pcps.h>
#include <src/input.h>
#include <src/wavefunction.h>
#include <src/accumulator.h>
#include <src/accumulator_energy.h>

namespace pcps {

  // forward declaration of thread accumulator
  template <class S, class PART, class REF> class GradientThreadAccum;

  //-----------------------------------------------------------------------------------
  // pcps::GradientProcessAccum -- Accumulates gradient data for one process and
  //                               applies a steepest descent step to the wavefunction.
  //-----------------------------------------------------------------------------------

  template <class S, class PART, class REF> class GradientProcessAccum : public pcps::ProcessAccum<S, PART, REF> {

    private:

      // data members
      const pcps::OptSettings _os;                // settings for wavefunction optimization
      int _n_corr_der;                            // number of correlator derivative elements on this process
      int _max_n_corr_der;                        // number of correlator derivative elements on the process with the most
      std::vector<int>    _corr_der_indices;      // indices for correlator derivatives computed by this processes
      std::vector<S> _corr_der_values;       // values for correlator derivatives computed by this processes
      std::vector<int>    _all_corr_der_indices;  // indices for correlator derivatives from all processes
      std::vector<S> _all_corr_der_values;   // values for correlator derivatives from all processes
      std::vector<int>    _all_n_corr_der;        // number of correlator derivative elements on each process
      std::vector<S> _agp_grad;              // gradient with respect to agp weight variables
      std::vector<S> _orb_grad;              // gradient with respect to restricted orbital coefficient variables
      std::vector<S> _par_grad;              // gradient with respect to pairing matrix variables

    public:

      // give the corresponding thread accumulator access to private and protected members
      friend class pcps::GradientThreadAccum<S, PART, REF>;

      // constructor
      GradientProcessAccum(const pcps::Input & userinp, const pcps::Wavefunction<S> & wfn, const long int sample_length)
        : pcps::ProcessAccum<S, PART, REF>(userinp, sample_length),
          _os(userinp, wfn)
      {

        // not yet implemented for translationally invariant correlators
        for (typename std::vector< pcps::Correlator<S> >::const_iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++)
          if (corr->ti())
            throw pcps::Exception("gradient accumulators not implemented for translationally invariant correlators");

        // initialize vectors
        const size_t sample_per_proc = size_t(this->_sample_length) * size_t(this->_nthread);
        const S z = pcps::zero<S>();
        if (_os.opt_cor) {
          _corr_der_indices.assign(     size_t(_os.n_var_corr) * sample_per_proc                        , 0);
          _corr_der_values.assign(      size_t(_os.n_var_corr) * sample_per_proc                        , z);
          _all_corr_der_indices.assign( size_t(_os.n_var_corr) * sample_per_proc * size_t(this->_nproc) , 0);
          _all_corr_der_values.assign(  size_t(_os.n_var_corr) * sample_per_proc * size_t(this->_nproc) , z);
          _all_n_corr_der.assign(       size_t(this->_nproc)                                            , 0);
        }
        if (_os.opt_agp) _agp_grad.assign(size_t(_os.n_agp_var), z);
        if (_os.opt_orb) _orb_grad.assign(size_t(_os.n_orb_var), z);
        if (_os.opt_par) _par_grad.assign(size_t(_os.n_par_var), z);

        // reset data
        this->reset();

      }

      // function to create a thread accumulator linked to this process accumulator
      pcps::ThreadAccum<S, PART, REF> * create_thread_accum(const int tid, pcps::Arena & arena);

      // function to reset the accumulator
      void reset() {
        if (_os.opt_agp) pcps::xscal(_os.n_agp_var, pcps::zero<S>(), &_agp_grad[0], 1);
        if (_os.opt_orb) pcps::xscal(_os.n_orb_var, pcps::zero<S>(), &_orb_grad[0], 1);
        if (_os.opt_par) pcps::xscal(_os.n_par_var, pcps::zero<S>(), &_par_grad[0], 1);
        if (_os.opt_cor) {
          _n_corr_der = 0;
          _max_n_corr_der = 0;
          for (int i = 0; i < _all_n_corr_der.size(); i++)
            _all_n_corr_der[i] = 0;
        }
      }

      // function to add the gradient, scaled by a step size, to the wavefunction's variables
      void update_wfn(const pcps::Input & userinp, pcps::Wavefunction<S> & wfn) {

        // do not add the gradient if the norm is too large
        if (this->get_grad_norm() > userinp.max_grad_norm()) return;

        // add agp weight, orbital coefficient, and pairing matrix gradients
        if (_os.opt_agp) pcps::xaxpy(_os.n_agp_var, -_os.step_size * pcps::unity<S>(), &_agp_grad[0], 1, &wfn.agp_weights()[0],           1);
        if (_os.opt_orb) pcps::xaxpy(_os.n_orb_var, -_os.step_size * pcps::unity<S>(), &_orb_grad[0], 1, &wfn.restricted_orb_coeffs()[0], 1);
        if (_os.opt_par) pcps::xaxpy(_os.n_par_var, -_os.step_size * pcps::unity<S>(), &_par_grad[0], 1, &wfn.pairing_matrix()[0],        1);

        // add correlator gradient
        if (_os.opt_cor) {
          int base = 0;
          for (int p = 0; p < this->_nproc; p++) {
            assert( _all_n_corr_der[p] % _os.n_var_corr == 0 );
            int x = base;
            const int n_accept = _all_n_corr_der[p] / _os.n_var_corr;
            for (int i = 0; i < n_accept; i++) {
              for (typename std::vector< pcps::Correlator<S> >::iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++) {
                if (corr->fixed()) continue;
                (*corr)[_all_corr_der_indices[x]] += -_os.step_size * _all_corr_der_values[x];
                x++;
              }
            }
            base += _max_n_corr_der;
          }
        }

        // if requested, force the pairing matrix to be symmetric
        if (_os.opt_par && userinp.symmetric_pairing() && userinp.ref_type() == pcps::keyword::Ref_Pairing)
          wfn.symmetrize_pairing_matrix();

      }

      // computes a norm for the gradient
      double get_grad_norm() const {

        // initialize the norm
        double retval = 0.0;

        // agp weight, orbital coefficient, and pairing matrix gradients
        if (_os.opt_agp) retval += pcps::real(pcps::xdotc(_os.n_agp_var, &_agp_grad[0], 1, &_agp_grad[0], 1));
        if (_os.opt_orb) retval += pcps::real(pcps::xdotc(_os.n_orb_var, &_orb_grad[0], 1, &_orb_grad[0], 1));
        if (_os.opt_par) retval += pcps::real(pcps::xdotc(_os.n_par_var, &_par_grad[0], 1, &_par_grad[0], 1));

        // correlator gradient
        if (_os.opt_cor) {
          int base = 0;
          for (int p = 0; p < this->_nproc; p++) {
            assert( _all_n_corr_der[p] % _os.n_var_corr == 0 );
            for (int i = base; i < base + _all_n_corr_der[p]; i++)
              retval += pcps::square_norm(_all_corr_der_values[i]);
            base += _max_n_corr_der;
          }
        }

        // return the gradient norm
        return std::abs(std::sqrt(retval));

      }

  };

  //--------------------------------------------------------------------------------
  // pcps::GradientThreadAccum -- Accumulates gradient data for one thread.
  //--------------------------------------------------------------------------------

  template <class S, class PART, class REF> class GradientThreadAccum : public pcps::ThreadAccum<S, PART, REF> {

    private:

      // pointer to the corresponding process accumulator
      pcps::GradientProcessAccum<S, PART, REF> * const _gpacc;

      // data members
      const pcps::OptSettings _os; // settings for wavefunction optimization
      int _n_accept;               // the number of configurations visited during sampling
      int    * _corr_der_indices;  // array for holding this clone's correlator gradient indices
      S * _corr_der_values;   // array for building this clone's correlator gradient values
      S * _local_energies;    // array for storing the local energies from the sample
      S * _agp_der_0;         // array used in building the agp weight gradient
      S * _agp_der_1;         // array used in building the agp weight gradient
      S * _agp_der_2;         // array used in building the agp weight gradient
      S * _orb_der_0;         // array used in building the orbital coefficient gradient
      S * _orb_der_1;         // array used in building the orbital coefficient gradient
      S * _orb_der_2;         // array used in building the orbital coefficient gradient
      S * _par_der_0;         // array used in building the pairing matrix gradient
      S * _par_der_1;         // array used in building the pairing matrix gradient
      S * _par_der_2;         // array used in building the pairing matrix gradient

    public:

      // constructor
      GradientThreadAccum(pcps::GradientProcessAccum<S, PART, REF> * const gpacc, const int tid, pcps::Arena & arena)
        : pcps::ThreadAccum<S, PART, REF>(gpacc, tid),
          _gpacc(gpacc),
          _os(gpacc->_os)
      {

        // initialize arrays
        const S z = pcps::zero<S>();
        _local_energies                    = arena.allocate_array<S>( size_t(this->_pacc->sample_length())                          , z);
        if (_os.opt_cor) _corr_der_indices = arena.allocate_array<   int>( size_t(this->_pacc->sample_length()) * size_t(_os.n_var_corr) , 0);
        if (_os.opt_cor) _corr_der_values  = arena.allocate_array<S>( size_t(this->_pacc->sample_length()) * size_t(_os.n_var_corr) , z);
        if (_os.opt_agp) _agp_der_0        = arena.allocate_array<S>( size_t(_os.n_agp_var)                                         , z);
        if (_os.opt_agp) _agp_der_1        = arena.allocate_array<S>( size_t(_os.n_agp_var)                                         , z);
        if (_os.opt_agp) _agp_der_2        = arena.allocate_array<S>( size_t(_os.n_agp_var)                                         , z);
        if (_os.opt_orb) _orb_der_0        = arena.allocate_array<S>( size_t(_os.n_orb_var)                                         , z);
        if (_os.opt_orb) _orb_der_1        = arena.allocate_array<S>( size_t(_os.n_orb_var)                                         , z);
        if (_os.opt_orb) _orb_der_2        = arena.allocate_array<S>( size_t(_os.n_orb_var)                                         , z);
        if (_os.opt_par) _par_der_0        = arena.allocate_array<S>( size_t(_os.n_par_var)                                         , z);
        if (_os.opt_par) _par_der_1        = arena.allocate_array<S>( size_t(_os.n_par_var)                                         , z);
        if (_os.opt_par) _par_der_2        = arena.allocate_array<S>( size_t(_os.n_par_var)                                         , z);

        // initialize data
        this->reset();

      }

      // function to reset the accumulator
      void reset() {
        _n_accept = 0;
        if (_os.opt_agp) pcps::xscal(_os.n_agp_var, pcps::zero<S>(), _agp_der_0, 1);
        if (_os.opt_agp) pcps::xscal(_os.n_agp_var, pcps::zero<S>(), _agp_der_1, 1);
        if (_os.opt_agp) pcps::xscal(_os.n_agp_var, pcps::zero<S>(), _agp_der_2, 1);
        if (_os.opt_orb) pcps::xscal(_os.n_orb_var, pcps::zero<S>(), _orb_der_0, 1);
        if (_os.opt_orb) pcps::xscal(_os.n_orb_var, pcps::zero<S>(), _orb_der_1, 1);
        if (_os.opt_orb) pcps::xscal(_os.n_orb_var, pcps::zero<S>(), _orb_der_2, 1);
        if (_os.opt_par) pcps::xscal(_os.n_par_var, pcps::zero<S>(), _par_der_0, 1);
        if (_os.opt_par) pcps::xscal(_os.n_par_var, pcps::zero<S>(), _par_der_1, 1);
        if (_os.opt_par) pcps::xscal(_os.n_par_var, pcps::zero<S>(), _par_der_2, 1);
      }

      // function to accumulate gradient information
      void accumulate(const long int iter, 
                      const long int count,
                      const bool accept,
                      const bool finished,
                      const S coeff_ratio,
                      const int * const new_config,
                      const int * const old_config,
                      pcps::CoeffValueFunc<S, PART, REF> & cvf,
                      pcps::ThreadAccum<S, PART, REF> ** accumulators,
                      const int n_accumulators) {

        if (accept || finished) {

          // get local energy
          {
            int i;
            for (i = 0; i < n_accumulators; i++)
              if (typeid(*accumulators[i]) == typeid(pcps::EnergyThreadAccum<S, PART, REF>)) {
                _local_energies[_n_accept] = ((pcps::EnergyThreadAccum<S, PART, REF> *)accumulators[i])->local_energy();
                break;
              }
            if (i == n_accumulators)
              throw pcps::Exception("error in pcps::GradientThreadAccum::accumulate.  Could not find an EnergyThreadAccum accumulator.");
          }

          // get derivative ratios
          cvf.get_compact_derivatives(old_config,
                                      _corr_der_indices + _n_accept * _os.n_var_corr,
                                      _corr_der_values  + _n_accept * _os.n_var_corr,
                                      _agp_der_0, _orb_der_0, _par_der_0);

          // prepare correlator gradient precursers
          if (_os.opt_cor) {
            pcps::xconj(_os.n_var_corr, _corr_der_values + _n_accept * _os.n_var_corr, 1);
            pcps::xscal(_os.n_var_corr, double(count) * pcps::unity<S>(), _corr_der_values + _n_accept * _os.n_var_corr, 1);
          }

          // prepare AGP weight gradient precursers
          if (_os.opt_agp) {
            pcps::xconj(_os.n_agp_var, _agp_der_0, 1);
            pcps::xaxpy(_os.n_agp_var, double(count) * _local_energies[_n_accept], _agp_der_0, 1, _agp_der_1, 1);
            pcps::xaxpy(_os.n_agp_var, double(count) *      pcps::unity<S>(), _agp_der_0, 1, _agp_der_2, 1);
          }

          // prepare orbital coefficient gradient precursers
          if (_os.opt_orb) {
            pcps::xconj(_os.n_orb_var, _orb_der_0, 1);
            pcps::xaxpy(_os.n_orb_var, double(count) * _local_energies[_n_accept], _orb_der_0, 1, _orb_der_1, 1);
            pcps::xaxpy(_os.n_orb_var, double(count) *      pcps::unity<S>(), _orb_der_0, 1, _orb_der_2, 1);
          }

          // prepare pairing matrix gradient precursers
          if (_os.opt_par) {
            pcps::xconj(_os.n_par_var, _par_der_0, 1);
            pcps::xaxpy(_os.n_par_var, double(count) * _local_energies[_n_accept], _par_der_0, 1, _par_der_1, 1);
            pcps::xaxpy(_os.n_par_var, double(count) *      pcps::unity<S>(), _par_der_0, 1, _par_der_2, 1);
          }

          // increment the number of accepted configurations
          _n_accept++;

        }

        if (finished) {

          // get the total energy
          S total_energy;
          {
            int i;
            for (i = 0; i < n_accumulators; i++)
              if (typeid(*accumulators[i]) == typeid(pcps::EnergyThreadAccum<S, PART, REF>)) {
                total_energy = ((pcps::EnergyThreadAccum<S, PART, REF> *)accumulators[i])->total_energy();
                break;
              }
            if (i == n_accumulators)
              throw pcps::Exception("error in pcps::GradientThreadAccum::accumulate.  Could not find an EnergyThreadAccum accumulator.");
          }

          // get the total number of samples
          const double total_samples = this->_pacc->total_samples();
          
          // compute gradient for correlator elements
          if (_os.opt_cor) {
            for (int i = 0; i < _n_accept; i++)
              pcps::xscal(_os.n_var_corr, 2.0 * (_local_energies[i] - total_energy) / total_samples, _corr_der_values + i * _os.n_var_corr, 1);
          }

          // compute gradient for agp weights
          if (_os.opt_agp) {
            pcps::xaxpy(_os.n_agp_var, -total_energy, _agp_der_2, 1, _agp_der_1, 1);
            pcps::xscal(_os.n_agp_var, 2.0 * pcps::unity<S>() / total_samples, _agp_der_1, 1);
          }

          // compute gradient for orbital coefficients
          if (_os.opt_orb) {
            pcps::xaxpy(_os.n_orb_var, -total_energy, _orb_der_2, 1, _orb_der_1, 1);
            pcps::xscal(_os.n_orb_var, 2.0 * pcps::unity<S>() / total_samples, _orb_der_1, 1);
          }

          // compute gradient for pairing matrix
          if (_os.opt_par) {
            pcps::xaxpy(_os.n_par_var, -total_energy, _par_der_2, 1, _par_der_1, 1);
            pcps::xscal(_os.n_par_var, 2.0 * pcps::unity<S>() / total_samples, _par_der_1, 1);
          }

          // add this thread's gradient to the process accumulator's gradient
          pthread_mutex_lock(this->_pacc->mutex());
          if (_os.opt_cor) {
            std::memcpy( &_gpacc->_corr_der_indices[_gpacc->_n_corr_der], _corr_der_indices, size_t(_n_accept) * size_t(_os.n_var_corr) * sizeof(int)    );
            std::memcpy( &_gpacc->_corr_der_values [_gpacc->_n_corr_der], _corr_der_values , size_t(_n_accept) * size_t(_os.n_var_corr) * sizeof(S) );
            _gpacc->_n_corr_der += _n_accept * _os.n_var_corr;
          }
          if (_os.opt_agp) pcps::xaxpy(_os.n_agp_var, pcps::unity<S>(), _agp_der_1, 1, &_gpacc->_agp_grad[0], 1);
          if (_os.opt_orb) pcps::xaxpy(_os.n_orb_var, pcps::unity<S>(), _orb_der_1, 1, &_gpacc->_orb_grad[0], 1);
          if (_os.opt_par) pcps::xaxpy(_os.n_par_var, pcps::unity<S>(), _par_der_1, 1, &_gpacc->_par_grad[0], 1);
          pthread_mutex_unlock(this->_pacc->mutex());

          // wait for other threads to add their gradients
          pthread_barrier_wait(this->_pacc->barrier());

          // thread 0 handles communication between processes
          if (this->_tid == 0) {

            // get MPI info
            const MPI::Comm & comm = MPI::COMM_WORLD;
            const int nproc = comm.Get_size();
            const int myrank = comm.Get_rank();

            // gradient for correlator elements
            if (_os.opt_cor) {
              comm.Allgather(&_gpacc->_n_corr_der, 1, MPI::INT, &_gpacc->_all_n_corr_der[0], 1, MPI::INT);
              _gpacc->_max_n_corr_der = *std::max_element(_gpacc->_all_n_corr_der.begin(), _gpacc->_all_n_corr_der.end());
              if (_gpacc->_max_n_corr_der > 100000000)
                throw pcps::Exception("too many correlator derivatives, need to implement allgather to operator on smaller chuncks");
              comm.Allgather(&_gpacc->_corr_der_indices[0],     int(_gpacc->_max_n_corr_der), MPI::INT,
                             &_gpacc->_all_corr_der_indices[0], int(_gpacc->_max_n_corr_der), MPI::INT);
              if (typeid(S) == typeid(double))
                comm.Allgather(&_gpacc->_corr_der_values[0],     int(_gpacc->_max_n_corr_der), MPI::DOUBLE,
                               &_gpacc->_all_corr_der_values[0], int(_gpacc->_max_n_corr_der), MPI::DOUBLE);
              else if (typeid(S) == typeid(std::complex<double>))
                comm.Allgather(&_gpacc->_corr_der_values[0],     int(_gpacc->_max_n_corr_der), MPI::DOUBLE_COMPLEX,
                               &_gpacc->_all_corr_der_values[0], int(_gpacc->_max_n_corr_der), MPI::DOUBLE_COMPLEX);
              else
                throw pcps::Exception("unknown S type in GradientAccumulator");
            }

            // gradient for agp weights
            if (_os.opt_agp) {
              pcps::xscal(_os.n_agp_var, pcps::zero<S>(), _agp_der_0, 1);
              pcps::allreduce(comm, &_gpacc->_agp_grad[0], _agp_der_0, _os.n_agp_var, MPI::SUM);
              pcps::xcopy(_os.n_agp_var, _agp_der_0, 1, &_gpacc->_agp_grad[0], 1);
            }

            // gradient for orbital coefficients
            if (_os.opt_orb) {
              pcps::xscal(_os.n_orb_var, pcps::zero<S>(), _orb_der_0, 1);
              pcps::allreduce(comm, &_gpacc->_orb_grad[0], _orb_der_0, _os.n_orb_var, MPI::SUM);
              pcps::xcopy(_os.n_orb_var, _orb_der_0, 1, &_gpacc->_orb_grad[0], 1);
            }

            // gradient for pairing matrix
            if (_os.opt_par) {
              pcps::xscal(_os.n_par_var, pcps::zero<S>(), _par_der_0, 1);
              pcps::allreduce(comm, &_gpacc->_par_grad[0], _par_der_0, _os.n_par_var, MPI::SUM);
              pcps::xcopy(_os.n_par_var, _par_der_0, 1, &_gpacc->_par_grad[0], 1);
            }

          }

          // wait for thread 0 to finish communication
          pthread_barrier_wait(this->_pacc->barrier());

        }

      }

  };

  //------------------------------------------------------------------------------------------
  // pcps::GradientProcessAccum::create_thread_accum -- Function to create a thread accumulator
  //                                                    from a process accumulator.
  //------------------------------------------------------------------------------------------

  template <class S, class PART, class REF>
  pcps::ThreadAccum<S, PART, REF> * pcps::GradientProcessAccum<S, PART, REF>::create_thread_accum(const int tid, pcps::Arena & arena) {
    return new(arena) pcps::GradientThreadAccum<S, PART, REF>(this, tid, arena);
  }

} // end namespace pcps

#endif
